package com.marcowu.study.spb.tracing.zipkin;

import brave.Tracing;
import brave.opentracing.BraveTracer;
import io.opentracing.Scope;
import io.opentracing.SpanContext;
import io.opentracing.Tracer;
import io.opentracing.noop.NoopTracerFactory;
import io.opentracing.propagation.Format;
import io.opentracing.propagation.TextMapExtractAdapter;
import io.opentracing.tag.Tags;
import zipkin2.Span;
import zipkin2.reporter.AsyncReporter;
import zipkin2.reporter.okhttp3.OkHttpSender;

import javax.ws.rs.core.MultivaluedMap;
import java.util.HashMap;

public class TraceUtil {
  private TraceUtil(){}

  public static Tracer getZipkinTracer(String serviceName) {
    try {
      // Configure a reporter, which controls how often spans are sent
      //   (the dependency is io.zipkin.reporter2:zipkin-sender-okhttp3)
      final OkHttpSender sender = OkHttpSender.create("http://localhost:9411/api/v2/spans");
      final AsyncReporter<Span> spanReporter = AsyncReporter.create(sender);

      // Now, create a Brave tracing component with the service name you want to see in Zipkin.
      //   (the dependency is io.zipkin.brave:brave)
      final Tracing braveTracing =
              Tracing.newBuilder()
                      .localServiceName(serviceName)
                      .spanReporter(spanReporter)
                      .build();

      // use this to create an OpenTracing Tracer
      return BraveTracer.create(braveTracing);
    } catch (Exception e) {
      e.printStackTrace();
      return NoopTracerFactory.create();
    }
  }

  public static Scope startServerSpan(Tracer tracer, javax.ws.rs.core.HttpHeaders httpHeaders,
                                      String operationName) {
    // format the headers for extraction
    MultivaluedMap<String, String> rawHeaders = httpHeaders.getRequestHeaders();
    final HashMap<String, String> headers = new HashMap<String, String>();
    for (String key : rawHeaders.keySet()) {
      headers.put(key, rawHeaders.get(key).get(0));
    }

    Tracer.SpanBuilder spanBuilder;
    try {
      SpanContext parentSpan = tracer.extract(Format.Builtin.HTTP_HEADERS, new TextMapExtractAdapter(headers));
      if (parentSpan == null) {
        spanBuilder = tracer.buildSpan(operationName);
      } else {
        spanBuilder = tracer.buildSpan(operationName).asChildOf(parentSpan);
      }
    } catch (IllegalArgumentException e) {
      spanBuilder = tracer.buildSpan(operationName);
    }
    return spanBuilder.withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_SERVER).startActive(true);
  }

}
